package org.apache.shenyu.plugin.trace;

import org.apache.shenyu.common.dto.RuleData;
import org.apache.shenyu.common.dto.SelectorData;
import org.apache.shenyu.common.enums.PluginEnum;
import org.apache.shenyu.common.utils.UUIDUtils;
import org.apache.shenyu.plugin.api.ShenyuPluginChain;
import org.apache.shenyu.plugin.base.AbstractShenyuPlugin;
import org.apache.shenyu.plugin.logging.common.constant.GenericLoggingConstant;
import org.jetbrains.annotations.NotNull;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

/**
 * traceId插件
 */
public class TraceIdPlugin extends AbstractShenyuPlugin {

    public static final String PLUGIN_NAME = "traceId";

    public static final String CUSTOM_HEADER_TRACE_ID = PLUGIN_NAME;

    @Override
    protected Mono<Void> doExecute(ServerWebExchange exchange, ShenyuPluginChain chain, SelectorData selector,
                                   RuleData rule) {
        // 生成或获取 traceId
        String traceId = generateOrGetTraceId();

        // 将 traceId 存储到 exchange 的属性中
        exchange.getAttributes().put(GenericLoggingConstant.SHENYU_AGENT_TRACE_ID, traceId);
        // 将 traceId 设置到请求头
        exchange = this.setRequestHttpHeader(exchange, traceId);
        // 将 traceId 设置到响应头
        exchange.getResponse().getHeaders().set(CUSTOM_HEADER_TRACE_ID, traceId);

        return chain.execute(exchange);
    }

    @NotNull
    private ServerWebExchange setRequestHttpHeader(ServerWebExchange exchange, String traceId) {
        // 创建新的请求头
        ServerHttpRequest request = exchange.getRequest();
        HttpHeaders newHeaders = new HttpHeaders();
        newHeaders.putAll(request.getHeaders());
        newHeaders.add(CUSTOM_HEADER_TRACE_ID, traceId);

        // 构建新请求
        ServerHttpRequest newRequest = request.mutate()
                .headers(HttpHeaders::clear)
                .headers(httpHeaders -> httpHeaders.putAll(newHeaders))
                .build();
        // 继续执行插件链
        return exchange.mutate().request(newRequest).build();
    }

    @Override
    public int getOrder() {
        // 设置一个较高的优先级，确保在大多数插件之前执行
        return PluginEnum.GLOBAL.getCode() - 1;
    }

    @Override
    public String named() {
        return TraceIdPlugin.PLUGIN_NAME;
    }

    private String generateOrGetTraceId() {
        return UUIDUtils.getInstance().generateShortUuid();
    }
}
